import numpy as np
from mpe.core import World, Agent, Landmark
from mpe.scenario import BaseScenario

scalar = 1

class Scenario(BaseScenario):
    def make_world(self):
        world = World()
        # set any world properties first
        world.dim_c = 2
        num_good_agents = 5 # prey
        num_adversaries = 2 # predator
        num_agents = num_adversaries + num_good_agents
        num_landmarks = 0
        # add agents
        world.agents = [Agent() for i in range(num_agents)]
        for i, agent in enumerate(world.agents):
            agent.name = 'agent %d' % i
            agent.collide = True
            agent.silent = True
            agent.adversary = True if i < num_adversaries else False
            agent.size = 0.075 if agent.adversary else 0.05
            agent.accel = 3.0 if agent.adversary else 1.5
            agent.max_speed = 1.0 if agent.adversary else 0.3
        # add landmarks
        world.landmarks = [Landmark() for i in range(num_landmarks)]
        for i, landmark in enumerate(world.landmarks):
            landmark.name = 'landmark %d' % i
            landmark.collide = True
            landmark.movable = False
            landmark.size = 0.0
            landmark.boundary = False
        # make initial conditions
        self.reset_world(world)
        return world

    def reset_world(self, world):
        # random properties for agents
        for i, agent in enumerate(world.agents):
            agent.color = np.array([0.35, 0.85, 0.35]) if not agent.adversary else np.array([0.85, 0.35, 0.35])
            # random properties for landmarks
        for i, landmark in enumerate(world.landmarks):
            landmark.color = np.array([0.25, 0.25, 0.25])
        # set random initial states
        prey_idx = 0
        for agent in world.agents:
            if agent.adversary:
                agent.state.p_pos = np.random.uniform(-0.1, +0.1, world.dim_p)

            else: # prey
                if prey_idx == 0:
                    agent.state.p_pos = np.array([1., 0.]) * scalar
                elif prey_idx == 1:
                    agent.state.p_pos = np.array([-1., 0.]) * scalar
                elif prey_idx == 2:
                    agent.state.p_pos = np.array([0.8, 0.6]) * scalar
                elif prey_idx == 3:
                    agent.state.p_pos = np.array([-0.8, 0.6]) * scalar
                elif prey_idx == 4:
                    agent.state.p_pos = np.array([0., -1.]) * scalar
                prey_idx += 1

            agent.state.p_vel = np.zeros(world.dim_p)
            agent.state.c = np.zeros(world.dim_c)
            agent.alive = True
        for i, landmark in enumerate(world.landmarks):
            if not landmark.boundary:
                landmark.state.p_pos = np.array([5, 0])

        self.caught_prey_set = set()
        self.caught_prey_level_sum = 0
        
        self.predators_position = {}
        self.prey_position = {}

    def is_collision(self, entity1, entity2):
        delta_pos = entity1.state.p_pos - entity2.state.p_pos
        dist = np.sqrt(np.sum(np.square(delta_pos)))
        dist_min = entity1.size + entity2.size
        return True if dist < dist_min * 2 else False

    # return all agents that are not adversaries
    def get_prey(self, world):
        return [agent for agent in world.agents if not agent.adversary]
    
    def get_prey_level(self, prey):
        return 2 if "S" in prey else 1
    
    def get_num_predator_nearby(self, prey):
        all_predators = self.get_predator(self.world)
        all_prey = self.get_prey(self.world)
        if prey == "S1":
            prey_entity = all_prey[0]
        elif prey == "S2":
            prey_entity = all_prey[1]
        elif prey == "R1":
            prey_entity = all_prey[2]
        elif prey == "R2":
            prey_entity = all_prey[3]
        elif prey == "R3":
            prey_entity = all_prey[4]
        return sum([self.is_collision(prey_entity, adv) for adv in all_predators])

    # return all adversarial agents
    def get_predator(self, world):
        return [agent for agent in world.agents if agent.adversary]
    
    def entity_distance(self, entity1, entity2):
        if len(entity1) == 1:
            entity1_pos = self.predators_position[entity1]
        elif len(entity1) == 2:
            entity1_pos = self.prey_position[entity1]
        if len(entity2) == 1:
            entity2_pos = self.predators_position[entity2]
        elif len(entity2) == 2:
            entity2_pos = self.prey_position[entity2]
        distance = np.sqrt(np.sum(np.square(entity1_pos - entity2_pos)))
        return distance
    
    def get_prey_name(self, prey_idx):
        if prey_idx == 0:
            k = "S1"
        elif prey_idx == 1:
            k = "S2"
        elif prey_idx == 2:
            k = "R1"
        elif prey_idx == 3:
            k = "R2"
        elif prey_idx == 4:
            k = "R3"
        return k
    
    def process_info(self, world):
        self.world = world
        predator_idx = 0
        prey_idx = 0
        for agent in world.agents:
            if agent.adversary:
                self.predators_position[str(predator_idx + 1)] = agent.state.p_pos
                predator_idx += 1
            else: # prey
                k = self.get_prey_name(prey_idx)
                self.prey_position[k] = agent.state.p_pos
                prey_idx += 1

    def reward(self, agent, world):
        
        self.process_info(world)
        
        rew = 0.0
        all_predators = self.get_predator(world)
        all_prey = self.get_prey(world)
        for prey_idx, prey in enumerate(all_prey):
            if not prey.alive:
                continue
            if prey_idx in [0, 1]: # stag
                prey_level = 2
            else: # rabbit
                prey_level = 1
            n_predator_nearby = sum([self.is_collision(prey, adv) for adv in all_predators])
            if n_predator_nearby >= prey_level:
                prey.alive = False
                prey.dead_pos = prey.state.p_pos
                prey.state.p_pos = np.array([999., 999.])
                rew += prey_level / 2
                self.caught_prey_set.add(self.get_prey_name(prey_idx))
                self.caught_prey_level_sum += prey_level
            elif 0 < n_predator_nearby < prey_level:
                rew -= 0.01
        return rew
    
    def get_dead_prey(self, world):
        return [prey_index for prey_index, prey in enumerate(self.get_prey(world)) if not prey.alive]

    def done(self, agent, world):
        return self.caught_prey_level_sum >= 2
        
    def llm0(self, agent, world):
        return 0
        
    def llm1(self, agent, world):
        return 0
        
    def llm2(self, agent, world):
        assert 0
        return 0
        
    def llm3(self, agent, world):
        assert 0
        return 0
        
    def llm4(self, agent, world):
        assert 0
        return 0
        
    def llm5(self, agent, world):
        assert 0
        return 0
        
    def llm6(self, agent, world):
        assert 0
        return 0

    def observation(self, agent, world):
        other_predator_pos = []
        for entity in world.agents:
            if not entity is agent and entity.adversary:
                other_predator_pos.append(entity.state.p_pos - agent.state.p_pos)
        other_prey_pos = []
        for entity in world.agents:
            if not entity is agent and not entity.adversary:
                if entity.alive:
                    other_prey_pos.append(entity.state.p_pos - agent.state.p_pos)
                else:
                    other_prey_pos.append(entity.dead_pos - agent.state.p_pos)
        obs = np.concatenate([agent.state.p_pos] + other_predator_pos + other_prey_pos)
        # for predator: self x, self y, tm x, tm y, prey 1 x, prey 1 y, ..., prey 5 x, prey 5 y
        # for prey: self x, self y, predator 1 x, predator 1 y, predator 2 x, predator 2 y, ...
        return obs
    
    